package cn.linkai.star.graph.core.service.impl;

import cn.hutool.core.date.DateTime;
import cn.linkai.star.graph.comfyui.client.pojo.ComfyuiTask;
import cn.linkai.star.graph.comfyui.client.pojo.MessageBase;
import cn.linkai.star.graph.core.job.RunTaskJob;
import cn.linkai.star.graph.core.service.ComfuiMessageService;
import cn.linkai.star.graph.core.service.RedisService;
import cn.linkai.star.graph.core.service.UserResultService;
import com.alibaba.fastjson2.JSON;
import lombok.RequiredArgsConstructor;
import org.redisson.api.RSemaphore;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;

@RequiredArgsConstructor
@Service
public class ComfuiMessageServiceImpl implements ComfuiMessageService {

    //注入
    private final RedisService redisService;
    private final WsNoticeServiceImpl wsNoticeService;
    private final UserResultService userResultService;
    private final RedissonClient redissonClient;

    @Override
    public void handleMessage(MessageBase messageBase) {
        if("progress".equals (messageBase.getType ())){
            // 处理进度
            handleProgress (messageBase);
        } else if("executed".equals (messageBase.getType ())) {
            // 处理生图结果
            handleExecuted (messageBase);
        } else if("status".equals (messageBase.getType ())){
            // 处理生图状态状态
            handleStatus (messageBase);
        }
    }

    private void handleStatus(MessageBase messageBase) {
        HashMap<String, Object> data = messageBase.getData ();
        HashMap<String, Object> status = (HashMap<String, Object>) data.get ("status");
        HashMap<String, Object> execInfonfo = (HashMap<String, Object>) status.get("exec_info");
        Integer queueRemaining = ( Integer) execInfonfo.get ("queue_remaining");

        //如果队列为0，则释放信号量
        if(queueRemaining == 0){
            RSemaphore semaphore = redissonClient.getSemaphore (RunTaskJob.TASK_RUN_SEMAPHORE);
            semaphore.release ();
            System.out.println (DateTime.now () + "\t\t释放许可证数量:"+semaphore.availablePermits ());

        }
    }

    /**
     * 处理生图结果
     * @param messageBase
     */
    private void handleExecuted(MessageBase messageBase) {
        //1、消息格式的转换
        HashMap<String, Object> data = messageBase.getData ();
        HashMap<String, Object> output = (HashMap<String, Object>)data.get("output");
        List<HashMap<String, Object>> images = (List<HashMap<String, Object>>) output.get ("images");
        List<String> urls = images.stream ().map (image ->
                String.format ("http://127.0.0.1:8188/view?filename=%s&type=%s&subfolder=%s", image.get ("filename"), image.get ("type"), image.get ("subfolder"))
        ).toList ();
        //2、redis中获取任务
        String promptId = (String) data.get ("prompt_id");
        ComfyuiTask task = redisService.getStartedTask (promptId);
        if(task == null){
            return;
        }

        //3、推送消息给前端
        HashMap<String, Object> result = new HashMap<> ();
        result.put ("type","imageResult");
        result.put ("urls",urls);
        wsNoticeService.sendToUser (task.getWsClientId (), JSON.toJSONString (result));

        //4、调用UserResultService保存图片
        userResultService.insertUserResult(task, urls);
    }


    /**
     * 处理生图进度
     * @param messageBase
     */
    private void handleProgress(MessageBase messageBase) {
        //1、消息格式的转换
        HashMap<String, Object> data = messageBase.getData ();
        data.put ("type","progress");

        //2、redis中获取任务
        String promptId = (String) data.get ("prompt_id");
        ComfyuiTask task = redisService.getStartedTask (promptId);
        if(task == null){
            return;
        }

        //3、推送消息给前端
        wsNoticeService.sendToUser (task.getWsClientId (), JSON.toJSONString (data));

    }
}
